from bpy.types import Operator
import bpy
import bmesh
import math
from mathutils import Matrix
import numpy
import sys
sys.path.extend(['', 'D:\\blender-2.92.0-windows64\\2.92\\python\\bin\\python37.zip', 'D:\\blender-2.92.0-windows64\\2.92\\python\\DLLs', 'D:\\blender-2.92.0-windows64\\2.92\\python\\lib', 'D:\\blender-2.92.0-windows64\\2.92\\python\\bin', 'C:\\Users\\tomye\\AppData\\Roaming\\Python\\Python37\\site-packages', 'D:\\blender-2.92.0-windows64\\2.92\\python', 'D:\\blender-2.92.0-windows64\\2.92\\python\\lib\\site-packages'])
import geopandas as gpd

def cross(a, b):
    return (a[1] * b[2] - a[2] * b[1], a[2] * b[0] - a[0] * b[2], a[0] * b[1] - a[1] * b[0])

def dot(a, b):
    return a[0] * b[0] + a[1] * b[1] + a[2] * b[2]

def normalize(a):
    a = numpy.array(a)
    return numpy.sqrt(numpy.sum(numpy.power(a, 2)))

def cal_rotate_matrix(a, b):
    rot_axis = cross(b, a)
    rot_angle = math.acos(dot(a, b) / normalize(a) / normalize(b))

    norm = normalize(rot_axis)
    rot_mat = numpy.zeros((4, 4), dtype = "float32")

    rot_axis = (rot_axis[0] / norm, rot_axis[1] / norm, rot_axis[2] / norm)

    rot_mat[0, 0] = math.cos(rot_angle) + rot_axis[0] * rot_axis[0] * (1 - math.cos(rot_angle))
    rot_mat[0, 1] = rot_axis[0] * rot_axis[1] * (1 - math.cos(rot_angle)) - rot_axis[2] * math.sin(rot_angle)
    rot_mat[0, 2] = rot_axis[1] * math.sin(rot_angle) + rot_axis[0] * rot_axis[2] * (1 - math.cos(rot_angle))
    rot_mat[0,3] = 0.0

    rot_mat[1, 0] = rot_axis[2] * math.sin(rot_angle) + rot_axis[0] * rot_axis[1] * (1 - math.cos(rot_angle))
    rot_mat[1, 1] = math.cos(rot_angle) + rot_axis[1] * rot_axis[1] * (1 - math.cos(rot_angle))
    rot_mat[1, 2] = -rot_axis[0] * math.sin(rot_angle) + rot_axis[1] * rot_axis[2] * (1 - math.cos(rot_angle))
    rot_mat[1,3] = 0.0

    rot_mat[2, 0] = -rot_axis[1] * math.sin(rot_angle) + rot_axis[0] * rot_axis[2] * (1 - math.cos(rot_angle))
    rot_mat[2, 1] = rot_axis[0] * math.sin(rot_angle) + rot_axis[1] * rot_axis[2] * (1 - math.cos(rot_angle))
    rot_mat[2, 2] = math.cos(rot_angle) + rot_axis[2] * rot_axis[2] * (1 - math.cos(rot_angle))
    rot_mat[2,3] = 0.0
    
    rot_mat[3,0] = rot_mat[3,1] = rot_mat[3,2] = 0.0
    rot_mat[3,3] = 1.0

    return Matrix(rot_mat)


# ------------------------------------------------------------------------------
# Utility Functions

def set_smooth(obj):
    """ Enable smooth shading on an mesh object """

    for face in obj.data.polygons:
        face.use_smooth = True


def object_from_data(data, name, scene, select=True):
    """ Create a mesh object and link it to a scene """

    mesh = bpy.data.meshes.new(name)
    mesh.from_pydata(data['verts'], data['edges'], data['faces'])

    obj = bpy.data.objects.new(name, mesh)
    scene.collection.objects.link(obj)

    # bpy.context.view_layer.objects.active = obj
    # obj.select_set(True)

    mesh.update(calc_edges=True)
    mesh.validate(verbose=True)

    return obj


def recalculate_normals(mesh):
    """ Make normals consistent for mesh """

    bm = bmesh.new()
    bm.from_mesh(mesh)

    bmesh.ops.recalc_face_normals(bm, faces=bm.faces)

    bm.to_mesh(mesh)
    bm.free()


# ------------------------------------------------------------------------------
# Geometry functions

def vertex_circle(segments, z):
    """ Return a ring of vertices """
    verts = []

    for i in range(segments):
        angle = (math.pi*2) * i / segments
        verts.append((math.cos(angle), math.sin(angle), z))

    return verts


def face(segments, i, row):
    """ Return a face on a cylinder """

    if i == segments - 1:
        ring_start = segments * row
        base = segments * (row + 1)

        return (base - 1, ring_start, base, (base + segments) - 1)

    else:
        base = (segments * row) + i
        return (base, base + 1, base + segments + 1, base + segments)


def bottom_cap(verts, faces, segments, cap='NGON'):
    """ Build bottom caps as triangle fans """

    if cap == 'TRI':
        verts.append((0, 0, 0))
        center_vert = len(verts) - 1

        [faces.append((i, i+1, center_vert)) for i in range(segments - 1)]
        faces.append((segments - 1, 0, center_vert))

    elif cap == 'NGON':
        faces.append([i for i in range(segments)])

    else:
        print('[!] Passed wrong type to bottom cap')


def top_cap(verts, faces, segments, rows, cap='NGON'):
    """ Build top caps as triangle fans """

    if cap == 'TRI':
        verts.append((0, 0, rows - 1))
        center_vert = len(verts) - 1
        base = segments * (rows - 1)

        [faces.append((base+i, base+i+1, center_vert))
                       for i in range(segments - 1)]

        faces.append((segments * rows - 1, base, center_vert))

    elif cap == 'NGON':
        base = (rows - 1) * segments
        faces.append([i + base for i in range(segments)])

    else:
        print('[!] Passed wrong type to top cap')


# ------------------------------------------------------------------------------
# Main Functions

def make_circle(name, segments=32, fill=None):
    """ Make a circle """

    data = {
            'verts': vertex_circle(segments, 0),
            'edges': [],
            'faces': [],
           }

    if fill:
        bottom_cap(data['verts'], data['faces'], segments, fill)
    else:
        data['edges'] = [(i, i+1) for i in range(segments)]
        data['edges'].append((segments - 1, 0))

    scene = bpy.context.scene
    return object_from_data(data, name, scene)


def make_cylinder(name, segments=16, rows=4, cap=None):
    """ Make a cylinder """

    data = { 'verts': [], 'edges': [], 'faces': [] }

    for z in range(rows):
        data['verts'].extend(vertex_circle(segments, z))

    for i in range(segments):
        for row in range(0, rows - 1):
            data['faces'].append(face(segments, i, row))

    if cap:
        bottom_cap(data['verts'], data['faces'], segments, cap)
        top_cap(data['verts'], data['faces'], segments, rows, cap)


    scene = bpy.context.scene
    obj = object_from_data(data, name, scene)
    recalculate_normals(obj.data)
    set_smooth(obj)

    # bevel = obj.modifiers.new('Bevel', 'BEVEL')
    # bevel.limit_method = 'ANGLE'

    # obj.modifiers.new('Edge Split', 'EDGE_SPLIT')

    return obj

def draw_pipe(name,x1,y1,z1,x2,y2,z2,diameter,offsetX,offsetY,offsetZ):
    # (x1,y1,z1)->(x2,y2,z2) 管径 diameter
    x = x2-x1
    y = y2-y1
    z = z2-z1
    # make_circle('Circle', 64)
    obj = make_cylinder(name, 16, 2) # 画一个两层 16圆节点的管道
    # 管径0.2
    obj.data.transform(Matrix.Scale(diameter, 4, (1,0,0)))
    obj.data.transform(Matrix.Scale(diameter, 4, (0,1,0)))
    # 管道长10
    length = math.sqrt(x**2+y**2+z**2)
    obj.data.transform(Matrix.Scale(length, 4, (0,0,1)))

    a = (0,0,1)
    b = (x,y,z)
    rot_mat = cal_rotate_matrix(b,a)
    obj.matrix_world @= rot_mat
    offset = (x1-offsetX,y1-offsetY,z1-offsetZ)
    obj.location = [i for i in offset]
    


class Draw_Cylinder_Op(Operator):
    bl_idname = "object.draw_cylinder"
    bl_label = "draw cylinder"
    bl_description = "draw a cylinder"
    
    def execute(self,context):
        data = gpd.read_file(r"C:\Users\tomye\Documents\中线.geojson")
        offsetX = data.geometry[0].centroid.x
        offsetY = data.geometry[0].centroid.y
        offsetZ = 0.0
        print(data.geometry.centroid)
        names = []
        for index,row in data.iterrows():
            name =row['管道编码']
            x1 = row.geometry.coords.xy[0][0]
            y1 = row.geometry.coords.xy[1][0]
            z1 = row['起点高程']
            x2 = row.geometry.coords.xy[0][1]
            y2 = row.geometry.coords.xy[1][1]
            z2 = row['终点高程']
            diameter = row['管径']/1000
            if name in names:
                print(name + "已经存在")
            else:
                names.append(name)
                draw_pipe(name,x1,y1,z1,x2,y2,z2,diameter,offsetX,offsetY,offsetZ)

        
        
        print(offsetX)
        print(offsetY) 
        # draw_pipe('pipe1',2.0,2.0,2.0,15.0,1.0,19.0,0.2)

        return {'FINISHED'}